import UIKit
import Metal
import MetalKit

// UIImage和MTLTexture之间的转换扩展
extension UIImage {
    func toMetalTexture(device: MTLDevice) -> MTLTexture? {
        guard let cgImage = self.cgImage else { return nil }
        
        let width = cgImage.width
        let height = cgImage.height
        let bytesPerPixel = 4
        let bytesPerRow = bytesPerPixel * width
        let bitsPerComponent = 8
        
        let pixelBuffer = UnsafeMutablePointer<UInt8>.allocate(capacity: width * height * bytesPerPixel)
        defer { pixelBuffer.deallocate() }
        
        let colorSpace = CGColorSpaceCreateDeviceRGB()
        let bitmapInfo = CGImageAlphaInfo.premultipliedLast.rawValue
        
        guard let context = CGContext(
            data: pixelBuffer,
            width: width,
            height: height,
            bitsPerComponent: bitsPerComponent,
            bytesPerRow: bytesPerRow,
            space: colorSpace,
            bitmapInfo: bitmapInfo
        ) else { return nil }
        
        context.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))
        
        let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(
            pixelFormat: .bgra8Unorm,
            width: width,
            height: height,
            mipmapped: false
        )
        textureDescriptor.usage = .shaderRead
        
        guard let texture = device.makeTexture(descriptor: textureDescriptor) else { return nil }
        
        let region = MTLRegionMake2D(0, 0, width, height)
        texture.replace(region: region, mipmapLevel: 0, withBytes: pixelBuffer, bytesPerRow: bytesPerRow)
        
        return texture
    }
}

extension MTLTexture {
    var size: CGSize {
        CGSize(width: width, height: height)
    }
    
    func toUIImage() -> UIImage? {
        let width = self.width
        let height = self.height
        let bytesPerRow = width * 4
        
        let pixelBuffer = UnsafeMutablePointer<UInt8>.allocate(capacity: width * height * 4)
        defer { pixelBuffer.deallocate() }
        
        let region = MTLRegionMake2D(0, 0, width, height)
        self.getBytes(pixelBuffer, bytesPerRow: bytesPerRow, from: region, mipmapLevel: 0)
        
        let colorSpace = CGColorSpaceCreateDeviceRGB()
        let bitmapInfo = CGImageAlphaInfo.premultipliedLast.rawValue
        
        guard let context = CGContext(
            data: pixelBuffer,
            width: width,
            height: height,
            bitsPerComponent: 8,
            bytesPerRow: bytesPerRow,
            space: colorSpace,
            bitmapInfo: bitmapInfo
        ), let cgImage = context.makeImage() else {
            return nil
        }
        
        return UIImage(cgImage: cgImage)
    }
}
    
